{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import subprocess\n",
    "from ast import literal_eval\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import MDAnalysis as mda\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from MDAnalysis.analysis import rms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FrameFlow Version\n",
    "start_idx_dict = {\n",
    "    \"1prw\": [15, 51],\n",
    "    \"1bcf\": [90, 122, 46, 17],\n",
    "    \"5tpn\": [108],\n",
    "    \"3ixt\": [0],\n",
    "    \"4jhw\": [144, 37],\n",
    "    \"4zyp\": [357],\n",
    "    \"5wn9\": [1],\n",
    "    \"5ius\": [88, 34],\n",
    "    \"5yui\": [89, 114, 194],\n",
    "    \"6vw1\": [5, 45],\n",
    "    \"1qjg\": [37, 13, 98],\n",
    "    \"1ycr\": [2],\n",
    "    \"2kl8\": [0, 27],\n",
    "    \"7mrx\": [25],\n",
    "    \"5trv\": [45],\n",
    "    \"6e6r\": [22],\n",
    "    \"6exz\": [25],\n",
    "}\n",
    "end_idx_dict = {\n",
    "    \"1prw\": [34, 70],\n",
    "    \"1bcf\": [98, 129, 53, 24],\n",
    "    \"5tpn\": [126],\n",
    "    \"3ixt\": [23],\n",
    "    \"4jhw\": [159, 43],\n",
    "    \"4zyp\": [371],\n",
    "    \"5wn9\": [20],\n",
    "    \"5ius\": [109, 53],\n",
    "    \"5yui\": [93, 116, 196],\n",
    "    \"6vw1\": [23, 63],\n",
    "    \"1qjg\": [37, 13, 98],\n",
    "    \"1ycr\": [10],\n",
    "    \"2kl8\": [6, 78],\n",
    "    \"7mrx\": [46],\n",
    "    \"5trv\": [69],\n",
    "    \"6e6r\": [34],\n",
    "    \"6exz\": [39],\n",
    "}\n",
    "\n",
    "motif_name_mapping = {\n",
    "    \"1PRW\": \"1prw\",\n",
    "    \"1BCF\": \"1bcf\",\n",
    "    \"5TPN\": \"5tpn\",\n",
    "    \"5IUS\": \"5ius\",\n",
    "    \"3IXT\": \"3ixt\",\n",
    "    \"5YUI\": \"5yui\",\n",
    "    \"1QJG\": \"1qjg\",\n",
    "    \"1YCR\": \"1ycr\",\n",
    "    \"2KL8\": \"2kl8\",\n",
    "    \"7MRX_60\": \"7mrx\",\n",
    "    \"7MRX_85\": \"7mrx\",\n",
    "    \"7MRX_128\": \"7mrx\",\n",
    "    \"4JHW\": \"4jhw\",\n",
    "    \"4ZYP\": \"4zyp\",\n",
    "    \"5WN9\": \"5wn9\",\n",
    "    \"5TRV_short\": \"5trv\",\n",
    "    \"5TRV_med\": \"5trv\",\n",
    "    \"5TRV_long\": \"5trv\",\n",
    "    \"6E6R_short\": \"6e6r\",\n",
    "    \"6E6R_med\": \"6e6r\",\n",
    "    \"6E6R_long\": \"6e6r\",\n",
    "    \"6EXZ_short\": \"6exz\",\n",
    "    \"6EXZ_med\": \"6exz\",\n",
    "    \"6EXZ_long\": \"6exz\",\n",
    "}\n",
    "\n",
    "import numpy as np\n",
    "from Bio import PDB\n",
    "\n",
    "\n",
    "def calculate_avg_plddt(pdb_file):\n",
    "    parser = PDB.PDBParser(QUIET=True)\n",
    "\n",
    "    structure = parser.get_structure(\"protein\", pdb_file)\n",
    "\n",
    "    plddt_values = []\n",
    "    for model in structure:\n",
    "        for chain in model:\n",
    "            for residue in chain:\n",
    "                if \"CA\" in residue:\n",
    "                    ca_atom = residue[\"CA\"]\n",
    "                    plddt = ca_atom.get_bfactor()\n",
    "                    plddt_values.append(plddt)\n",
    "\n",
    "    if plddt_values:\n",
    "        avg_plddt = np.mean(plddt_values)\n",
    "        return avg_plddt\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "\n",
    "\n",
    "def calc_rmsd_tmscore(\n",
    "    pdb_name,\n",
    "    reference_PDB,\n",
    "    scaffold_pdb_path=None,\n",
    "    scaffold_info_path=None,\n",
    "    ref_motif_starts=[30],\n",
    "    ref_motif_ends=[44],\n",
    "    output_path=None,\n",
    "):\n",
    "    \"Calculate RMSD between reference structure and generated structure over the defined motif regions\"\n",
    "\n",
    "    motif_df = pd.read_csv(\n",
    "        os.path.join(scaffold_info_path, f\"{pdb_name}.csv\"), index_col=0\n",
    "    )\n",
    "    results = []\n",
    "    for pdb in os.listdir(\n",
    "        scaffold_pdb_path\n",
    "    ):  # This needs to be in numerical order to match new_starts file\n",
    "        if not pdb.endswith(\".pdb\"):\n",
    "            continue\n",
    "        ref = mda.Universe(reference_PDB)\n",
    "        predict_PDB = os.path.join(scaffold_pdb_path, pdb)\n",
    "        u = mda.Universe(predict_PDB)\n",
    "\n",
    "        ref_selection = \"name CA and resnum \"\n",
    "        u_selection = \"name CA and resnum \"\n",
    "\n",
    "        i = int(pdb.split(\"_\")[1].split(\".\")[0])\n",
    "        new_motif_starts = literal_eval(motif_df[\"start_idxs\"].iloc[i])\n",
    "        new_motif_ends = literal_eval(motif_df[\"end_idxs\"].iloc[i])\n",
    "\n",
    "        ref_selection_list = []\n",
    "        for j in range(len(ref_motif_starts)):\n",
    "            ref_selection_temp = (\n",
    "                ref_selection\n",
    "                + str(ref_motif_starts[j])\n",
    "                + \":\"\n",
    "                + str(ref_motif_ends[j])\n",
    "                + \" \"\n",
    "            )\n",
    "            ref_selection_list.append(ref_selection_temp)\n",
    "            u_selection += (\n",
    "                str(new_motif_starts[j] + 1)\n",
    "                + \":\"\n",
    "                + str(new_motif_ends[j] + 1)\n",
    "                + \" \"\n",
    "            )\n",
    "\n",
    "        ref_ordered = ref.select_atoms(ref_selection_list[0])\n",
    "        for j in range(1, len(ref_selection_list)):\n",
    "            ref_ordered = ref_ordered + ref.select_atoms(ref_selection_list[j])\n",
    "\n",
    "        # print(\"U SELECTION\", u_selection)\n",
    "        # print(\"SEQUENCE\", i)\n",
    "        # print(\"ref\", ref_ordered.resnames)\n",
    "        # print(\"gen\", u.select_atoms(u_selection).resnames)\n",
    "        # This asserts that the motif sequences are the same - if you get this error something about your indices are incorrect - check chain/numbering\n",
    "        assert len(ref_ordered.resnames) == len(\n",
    "            u.select_atoms(u_selection).resnames\n",
    "        ), \"Motif \\\n",
    "                                                                    lengths do not match, check PDB preprocessing\\\n",
    "                                                                    for extra residues\"\n",
    "\n",
    "        assert (\n",
    "            ref_ordered.resnames == u.select_atoms(u_selection).resnames\n",
    "        ).all(), \"Resnames for\\\n",
    "                                                                    motifRMSD do not match, check indexing\"\n",
    "        rmsd = rms.rmsd(\n",
    "            u.select_atoms(u_selection).positions,\n",
    "            # coordinates to align\n",
    "            ref_ordered.positions,\n",
    "            # reference coordinates\n",
    "            center=True,  # subtract the center of geometry\n",
    "            superposition=True,\n",
    "        )  # superimpose coordinates\n",
    "\n",
    "        temp_file = open(os.path.join(output_path, \"temp_tmscores.txt\"), \"w\")\n",
    "\n",
    "        subprocess.call(\n",
    "            [\"./TMscore\", reference_PDB, predict_PDB, \"-seq\"], stdout=temp_file\n",
    "        )\n",
    "        with open(os.path.join(output_path, \"temp_tmscores.txt\"), \"r\") as f:\n",
    "            for line in f:\n",
    "                if len(line.split()) > 1 and \"TM-score\" == line.split()[0]:\n",
    "                    tm_score = line.split()[2]\n",
    "                    break\n",
    "\n",
    "        plddt = calculate_avg_plddt(predict_PDB)\n",
    "        results.append((pdb_name, i, rmsd, plddt, tm_score))\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DPLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate motif rmsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaffold_dir = \"../generation-results/dplm_650m/motif_scaffold\"\n",
    "motif_pdb_dir = \"../data-bin/scaffolding-pdbs\"\n",
    "\n",
    "output_dir = os.path.join(scaffold_dir, \"scaffold_results\")\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "results = []\n",
    "\n",
    "for ori_pdb, pdb in motif_name_mapping.items():\n",
    "    print(pdb)\n",
    "    ref_motif_starts = start_idx_dict[pdb]\n",
    "    ref_motif_ends = end_idx_dict[pdb]\n",
    "    reference_PDB = os.path.join(\n",
    "        \"../data-bin/scaffolding-pdbs\", pdb + \"_reference.pdb\"\n",
    "    )\n",
    "    with open(reference_PDB) as f:\n",
    "        line = f.readline()\n",
    "        ref_basenum = int(line.split()[5])\n",
    "    ref_motif_starts = [num + ref_basenum for num in ref_motif_starts]\n",
    "    ref_motif_ends = [num + ref_basenum for num in ref_motif_ends]\n",
    "    results += calc_rmsd_tmscore(\n",
    "        pdb_name=ori_pdb,\n",
    "        reference_PDB=reference_PDB,\n",
    "        scaffold_pdb_path=f\"{scaffold_dir}/scaffold_fasta/esmfold_pdb/{ori_pdb}\",\n",
    "        scaffold_info_path=f\"{scaffold_dir}/scaffold_info\",\n",
    "        ref_motif_starts=ref_motif_starts,\n",
    "        ref_motif_ends=ref_motif_ends,\n",
    "        output_path=output_dir,\n",
    "    )\n",
    "\n",
    "results = pd.DataFrame(\n",
    "    results, columns=[\"pdb_name\", \"index\", \"rmsd\", \"plddt\", \"tmscore\"]\n",
    ")\n",
    "results.to_csv(os.path.join(output_dir, \"rmsd_tmscore.csv\"), index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate success rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_success_scaffold(pdb):\n",
    "    total = len(pdb)\n",
    "    pdb[\"total\"] = total\n",
    "    pdb = pdb[(pdb[\"rmsd\"] < 1.0) & (pdb[\"plddt\"] > 70)]\n",
    "    return pdb\n",
    "\n",
    "\n",
    "rmsd_tmscore = pd.read_csv(os.path.join(output_dir, \"rmsd_tmscore.csv\"))\n",
    "success_scaffold = rmsd_tmscore.groupby(\"pdb_name\", as_index=False).apply(\n",
    "    cal_success_scaffold\n",
    ")\n",
    "success_scaffold_count = success_scaffold.groupby(\"pdb_name\").size()\n",
    "success_scaffold_count = success_scaffold_count.reset_index(\n",
    "    name=\"success_count\"\n",
    ")\n",
    "\n",
    "all_pdb = list(rmsd_tmscore[\"pdb_name\"].unique())\n",
    "success_pdb = list(success_scaffold_count[\"pdb_name\"])\n",
    "failed_pdb = list(set(all_pdb) - set(success_pdb))\n",
    "failed_scaffold_count = {\n",
    "    \"pdb_name\": failed_pdb,\n",
    "    \"success_count\": [0] * len(failed_pdb),\n",
    "}\n",
    "results = pd.concat(\n",
    "    [success_scaffold_count, pd.DataFrame(failed_scaffold_count)]\n",
    ").sort_values(\"pdb_name\")\n",
    "results.to_csv(os.path.join(output_dir, \"result.csv\"))\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DPLM-2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate motif rmsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaffold_dir = \"../generation-results/dplm2_650m/motif_scaffold\"\n",
    "motif_pdb_dir = \"../data-bin/scaffolding-pdbs\"\n",
    "\n",
    "output_dir = os.path.join(scaffold_dir, \"scaffold_results\")\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "results = []\n",
    "\n",
    "for ori_pdb, pdb in motif_name_mapping.items():\n",
    "    print(pdb)\n",
    "    ref_motif_starts = start_idx_dict[pdb]\n",
    "    ref_motif_ends = end_idx_dict[pdb]\n",
    "    reference_PDB = os.path.join(\n",
    "        \"../data-bin/scaffolding-pdbs\", pdb + \"_reference.pdb\"\n",
    "    )\n",
    "    with open(reference_PDB) as f:\n",
    "        line = f.readline()\n",
    "        ref_basenum = int(line.split()[5])\n",
    "    ref_motif_starts = [num + ref_basenum for num in ref_motif_starts]\n",
    "    ref_motif_ends = [num + ref_basenum for num in ref_motif_ends]\n",
    "    results += calc_rmsd_tmscore(\n",
    "        pdb_name=ori_pdb,\n",
    "        reference_PDB=reference_PDB,\n",
    "        scaffold_pdb_path=f\"{scaffold_dir}/scaffold_fasta/{ori_pdb}/esmfold_pdb\",\n",
    "        scaffold_info_path=f\"{scaffold_dir}/scaffold_info\",\n",
    "        ref_motif_starts=ref_motif_starts,\n",
    "        ref_motif_ends=ref_motif_ends,\n",
    "        output_path=output_dir,\n",
    "    )\n",
    "\n",
    "results = pd.DataFrame(\n",
    "    results, columns=[\"pdb_name\", \"index\", \"rmsd\", \"plddt\", \"tmscore\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_cat = []\n",
    "pdb_name_list = []\n",
    "for path in glob.glob(\n",
    "    scaffold_dir + \"/scaffold_fasta/**/**/**/eval/all_top_samples.csv\"\n",
    "):\n",
    "    df = pd.read_csv(path)\n",
    "    pdb_name = path.split(\"/\")[-3]\n",
    "    if pdb_name in pdb_name_list:\n",
    "        continue\n",
    "    pdb_name_list.append(pdb_name)\n",
    "    df[\"pdb_name\"] = pdb_name\n",
    "    df[\"index\"] = df[\"sample_path\"].apply(\n",
    "        lambda x: int(x.split(\"/\")[-2].split(\"_\")[1].split(\".\")[0])\n",
    "    )\n",
    "    df = df[[\"pdb_name\", \"index\", \"bb_rmsd\", \"bb_tmscore\"]]\n",
    "    df_cat.append(df)\n",
    "df_cat = pd.concat(df_cat, axis=0)\n",
    "results = pd.merge(results, df_cat, on=[\"pdb_name\", \"index\"])\n",
    "results.to_csv(os.path.join(output_dir, \"rmsd_tmscore.csv\"), index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate success rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_success_scaffold(pdb):\n",
    "    total = len(pdb)\n",
    "    pdb[\"total\"] = total\n",
    "    pdb = pdb[(pdb[\"rmsd\"] < 1.0) & (pdb[\"bb_tmscore\"] > 0.8)]\n",
    "    return pdb\n",
    "\n",
    "\n",
    "rmsd_tmscore = pd.read_csv(os.path.join(output_dir, \"rmsd_tmscore.csv\"))\n",
    "success_scaffold = rmsd_tmscore.groupby(\"pdb_name\", as_index=False).apply(\n",
    "    cal_success_scaffold\n",
    ")\n",
    "success_scaffold_count = success_scaffold.groupby(\"pdb_name\").size()\n",
    "success_scaffold_count = success_scaffold_count.reset_index(\n",
    "    name=\"success_count\"\n",
    ")\n",
    "\n",
    "all_pdb = list(rmsd_tmscore[\"pdb_name\"].unique())\n",
    "success_pdb = list(success_scaffold_count[\"pdb_name\"])\n",
    "failed_pdb = list(set(all_pdb) - set(success_pdb))\n",
    "failed_scaffold_count = {\n",
    "    \"pdb_name\": failed_pdb,\n",
    "    \"success_count\": [0] * len(failed_pdb),\n",
    "}\n",
    "results = pd.concat(\n",
    "    [success_scaffold_count, pd.DataFrame(failed_scaffold_count)]\n",
    ").sort_values(\"pdb_name\")\n",
    "results.to_csv(os.path.join(output_dir, \"result.csv\"))\n",
    "results"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dplm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
